from collections import defaultdict

from pandas.core.frame import DataFrame
from src.taskSignals import TaskSignals

from src.datahelper import StockDataHelper
import numpy as np


class BaseStrategy:
    def __init__(self) -> None:
        # 所有交易日
        self.trade_day_list = []
        # taskSigals 与主窗口连接的信号
        self.taskSignals: TaskSignals = None
        # 存储每次调仓日的决策
        self.stocks_log = defaultdict(int)
        # 各股仓位
        self.position = defaultdict(int)
        self.initial_money = 0
        self.daily_return = []
        # 各股仓位
        self.position = defaultdict(int)
        self.money = 0
        self.data_by_date = defaultdict(DataFrame)

    # 更新所有股票的总市值
    def update_portfolio_value(self, date):
        cur_date_series = self.data_by_date[date]['close']
        self.portfolio_value = 0
        for stock, cnt in self.position.items():
            p = cur_date_series.get(stock)
            if not np.isnan(p) and p:
                self.portfolio_value += p * cnt

    def log(self, msg):
        msg = str(msg)
        self.taskSignals.logSignal.emit(msg)

    def init_strategy(self):
        pass

    def updateBackTestingData(self):
        n = len(list(filter(lambda x: x > 0, self.daily_return)))

        # 更新回测结果
        self.taskSignals.updateBackTestingSignal.emit({
            'max_withdraw': StockDataHelper.max_drawdown(self.daily_return),
            'daily_return': self.daily_return,
            'final_money': self.money + self.portfolio_value,
            'profit_rate': (self.money + self.portfolio_value - self.initial_money) / self.initial_money * 100,
            'win_rate': n / len(self.daily_return) * 100,
            'win_lose_rate': n / (len(self.daily_return) - n),
            'trade_day_list': self.trade_day_list,
            'daily_return': self.daily_return
        })

    def sell_stocks(self, stocks: list, date):
        df = self.data_by_date[date]
        for stock in stocks:
            price = df.loc[stock, 'close']
            n = self.position[stock]
            if np.isnan(price) or n == 0:
                continue
            tot = n * price
            self.money += (1 - self.service_charge) * tot
            self.log(
                f'卖出{stock}, 卖出价{price}, 卖出{n}股, 当前资金{self.money}')
            self.position[stock] = 0

    # 等比例购买
    def buy_stocks(self, stocks: list, date, col='close'):
        if len(stocks) == 0:
            return
        df = self.data_by_date[date]
        each_money = self.money / len(stocks)
        for stock in stocks:
            price = df.loc[stock, col]
            if np.isnan(price):
                continue
            n = int(each_money // price)
            self.money -= n * price
            self.position[stock] += n
            self.log(f'买入{stock}, 买入价为{price}, 买入{n}股')

    def run(self):
        pass
